


from database import Database, OracleDB, DBType
import polars as pl

class ColumnsMgr:
    def __init__(self, db: Database):
        self.db = db
        self.cols_df = self.__query_all_columns_to_polars()
        self.uk_df = self.__query_unique_keys_to_polars()

    def get_uk_columns(self, owner: str, table_name: str):
        if self.uk_df.is_empty():
            raise Exception('Unique key not yet obtained')
        return self.uk_df.filter((pl.col('TABLE_NAME') == table_name.upper())& (pl.col('OWNER') == owner.upper())).get_column('COLUMN_NAME').to_list()

    def get_all_columns(self, owner: str, table_name: str):
        if self.cols_df.is_empty():
            raise Exception('Unique key not yet obtained')
        return self.cols_df.filter((pl.col('TABLE_NAME') == table_name.upper())& (pl.col('OWNER') == owner.upper())).get_column('COLUMN_NAME').to_list()
    
    def __query_all_columns_to_polars(self):
        try:
            if self.dbtype==DBType.ORACLE:
                # 定义 SQL 查询语句
                sql_query = """
        SELECT owner, table_name, column_name from all_tab_columns t1
        left join all_users t2
        on t1.owner = t2.username
        where t2.oracle_maintained = 'N'
                """
            elif self.dbtype==DBType.DMDB:
            # 定义 SQL 查询语句
            sql_query = """SELECT  owner, table_name, column_name 
                FROM 
                    DBA_TAB_COLUMNS t1
                LEFT JOIN 
                    DBA_USERS t2
                ON 
                    t1.owner = t2.username
                WHERE 
                    t2.username NOT LIKE 'SYS%' 
                    AND t2.username NOT LIKE 'SYSDBA'
                    AND t2.username NOT LIKE 'SYSAUDITOR'
                    AND t2.username NOT LIKE 'SYSSSO'
            """
            else:
                raise Exception('Unknown database type')

            results, columns = self.db.fetch_all(sql_query)

            # 创建 Polars DataFrame
            df = pl.DataFrame(results, orient='row', schema=columns)

            return df
        except Exception as e:
            print(f"查询出错: {e}")

    def __query_unique_keys_to_polars(self):
        try:
            # 定义 SQL 查询语句
            sql_query = """
    WITH
    -- 查询主键列
    pk_columns AS
    (SELECT c.owner, c.table_name, c.column_name, 1 AS priority
        FROM all_cons_columns c
        JOIN all_constraints a
          ON c.constraint_name = a.constraint_name
      WHERE a.constraint_type = 'P'),
    -- 查询唯一约束列（仅当无主键时）
    uk_columns AS
    (SELECT c.owner, c.table_name, c.column_name, 2 AS priority
        FROM all_cons_columns c
        JOIN all_constraints a
          ON c.constraint_name = a.constraint_name
      WHERE a.constraint_type = 'U'
        AND NOT EXISTS (SELECT 1 FROM pk_columns WHERE table_name = a.table_name)),
    -- 查询唯一索引列（仅当无主键和唯一约束时）
    uk_index_columns AS
    (SELECT i.owner, i.table_name, ic.column_name, 3 AS priority
        FROM all_ind_columns ic
        JOIN all_indexes i
          ON ic.index_name = i.index_name
      WHERE i.uniqueness = 'UNIQUE'
        AND NOT EXISTS (SELECT 1 FROM pk_columns WHERE table_name = i.table_name)
        AND NOT EXISTS (SELECT 1 FROM uk_columns WHERE table_name = i.table_name))
    -- 合并结果
    SELECT owner, table_name, column_name
      FROM (SELECT * FROM pk_columns UNION ALL SELECT * FROM uk_columns UNION ALL SELECT * FROM uk_index_columns)
      WHERE TABLE_NAME NOT LIKE 'BIN$%'
    ORDER BY priority, column_name
            """
            results, columns = self.db.fetch_all(sql_query)


            # 创建 Polars DataFrame
            df = pl.DataFrame(results, orient='row', schema=columns)


            return df
        except Exception as e:
            print(f"查询出错: {e}")


if __name__ == "__main__":
    db = OracleDB()
    db.connect()

    ukmgr = ColumnsMgr(db)
    cols = ukmgr.get_uk_columns('OWNER', 'TBL_XXX')
    
    print(cols)
    db.disconnect()